import asyncio
import uuid
from datetime import datetime

from typing import Dict, Optional, List, Any

from WebBackend.dataclass.TaskItem import TaskStatus
from WebBackend.utils import mylog


class Task:
    def __init__(
            self,
            command: str,
            task_id: Optional[str] = None,
            timeout: Optional[int] = None,
    ):
        self.task_id = task_id or str(uuid.uuid4())
        self.command = command
        self.status = TaskStatus.PENDING
        self.created_at = datetime.now()
        self.started_at: Optional[datetime] = None
        self.finished_at: Optional[datetime] = None
        self.timeout = timeout  # in seconds
        self.return_code: Optional[int] = None
        self.output: Optional[str] = None
        self.error: Optional[str] = None
        self._process: Optional[asyncio.subprocess.Process] = None

    async def execute(self):
        self.status = TaskStatus.RUNNING
        self.started_at = datetime.now()

        try:
            self._process = await asyncio.create_subprocess_shell(
                self.command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )

            try:
                stdout, stderr = await asyncio.wait_for(
                    self._process.communicate(),
                    timeout=self.timeout if self.timeout else None,
                )
                self.return_code = self._process.returncode
                self.output = stdout.decode("utf-8") if stdout else None
                self.error = stderr.decode("utf-8") if stderr else None

                if self.return_code == 0:
                    self.status = TaskStatus.COMPLETED
                    mylog.info(f"Task[{self.task_id}][{self.command}] executed finished")
                else:
                    self.status = TaskStatus.FAILED
                    mylog.warn(f"Task[{self.task_id}][{self.command}] executed failed")
            except asyncio.TimeoutError:
                self._process.kill()
                await self._process.wait()
                self.status = TaskStatus.FAILED
                self.error = f"Task timed out after {self.timeout} seconds"
                mylog.warn(f"Task[{self.task_id}][{self.command}] executed timeout")
        except Exception as e:
            self.status = TaskStatus.FAILED
            self.error = str(e)
        finally:
            self.finished_at = datetime.now()
            self._process = None

    async def cancel(self):
        if self._process and self._process.returncode is None:
            self._process.kill()
            await self._process.wait()
            self.status = TaskStatus.CANCELLED
            self.finished_at = datetime.now()
            return True
        return False

    def to_dict(self) -> Dict[str, Any]:
        return {
            "task_id": self.task_id,
            "command": self.command,
            "status": self.status.value,
            "created_at": self.created_at.isoformat(),
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "finished_at": self.finished_at.isoformat() if self.finished_at else None,
            "return_code": self.return_code,
            "output": self.output,
            "error": self.error,
            "timeout": self.timeout,
        }


class TaskStore:
    def __init__(self):
        self._tasks: Dict[str, Task] = {}
        self._lock = asyncio.Lock()

    async def get_task(self, task_id: str) -> Optional[Task]:
        """Retrieve a task by its ID."""
        async with self._lock:
            return self._tasks.get(task_id)

    async def list_tasks(self, status: Optional[TaskStatus] = None) -> List[Dict[str, Any]]:
        """List all tasks, optionally filtered by status."""
        async with self._lock:
            if status:
                return [task.to_dict() for task in self._tasks.values() if task.status == status]
            return [task.to_dict() for task in self._tasks.values()]

    async def create_task(self, command: str, timeout: Optional[int] = None, task_id: str = '') -> Task:
        """Create and store a new task."""
        async with self._lock:
            task = Task(command=command, timeout=timeout, task_id=task_id)
            self._tasks[task.task_id] = task
            return task

    async def cancel_task(self, task_id: str) -> bool:
        """Attempt to cancel a running task."""
        async with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                return False
            return await task.cancel()

    async def start_task(self, task_id: str):
        """Start a new task."""
        async with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                return False
            else:
                asyncio.create_task(task.execute())
                return True

    async def cleanup_completed_tasks(self, older_than_seconds: int = 86400) -> int:
        """Clean up completed tasks older than specified seconds (default: 24 hours)."""
        now = datetime.now()
        threshold = older_than_seconds
        removed = 0

        async with self._lock:
            for task_id in list(self._tasks.keys()):
                task = self._tasks[task_id]
                if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED):
                    if (now - task.finished_at).total_seconds() > threshold:
                        del self._tasks[task_id]
                        removed += 1
        return removed
